# -*- coding: utf-8 -*-
import json
from PySide6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QGridLayout, QFormLayout,
    QListWidget, QListWidgetItem, QTextEdit, QPushButton, QDialog, QLineEdit, QMessageBox)
from PySide6.QtGui import QValidator, QIntValidator, QRegularExpressionValidator
from triton_client_wrapper.triton_server_client import TritonServerClient

class ServerInfoDialog(QDialog):
    def __init__(self, parent = None):
        super(ServerInfoDialog, self).__init__(parent)
        self.parent_widget = parent
        self.server_info = {}
        dialog_layout = QVBoxLayout()
        # button layout
        confirm_button = QPushButton("确定")
        cancle_button = QPushButton("取消")
        button_layout = QHBoxLayout()
        confirm_button.clicked.connect(self.confirm_button_callback)
        cancle_button.clicked.connect(self.cancle_button_callback)
        button_layout.addWidget(confirm_button)
        button_layout.addWidget(cancle_button)

        # input lineedit layout
        input_layout = QFormLayout()
        self.ip_addr = QLineEdit()
        ip_re = "^((\d|[1-9]\d|1\d\d|2[0-4]\d|25[0-5])\.){3}(\d|[1-9]\d|1\d\d|2[0-4]\d|25[0-5])$"
        self.ip_validator = QRegularExpressionValidator(ip_re)
        self.ip_addr.setValidator(self.ip_validator)
        self.ip_addr.setPlaceholderText("0.0.0.0~255.255.255.255")
        self.port_validator = QIntValidator(0, 65535)
        self.http_port = QLineEdit()
        self.http_port.setValidator(self.port_validator)
        self.http_port.setPlaceholderText("0~65535")
        self.grpc_port = QLineEdit()
        self.grpc_port.setValidator(self.port_validator)
        self.grpc_port.setPlaceholderText("0~65535")
        self.server_name = QLineEdit()
        self.server_name.setPlaceholderText("不能够为空")
        input_layout.addRow("服务ip地址", self.ip_addr)
        input_layout.addRow("服务http端口", self.http_port)
        input_layout.addRow("服务grpc端口", self.grpc_port)
        input_layout.addRow("服务名称", self.server_name)
        dialog_layout.addLayout(input_layout)
        dialog_layout.addLayout(button_layout)
        self.setLayout(dialog_layout)

    def clear_dialog(self):
        self.ip_addr.setText("")
        self.http_port.setText("")
        self.grpc_port.setText("")
        self.server_name.setText("")

    def confirm_button_callback(self):
        self.server_info = {}
        ip_state, ip_addr, pos = self.ip_validator.validate(self.ip_addr.text(), 0)
        http_state, http_port, pos = self.port_validator.validate(self.http_port.text(), 0)
        grpc_state, grpc_port, pos = self.port_validator.validate(self.grpc_port.text(), 0)
        server_name = self.server_name.text()
        if QValidator.Acceptable.value != ip_state.value:
            QMessageBox.critical(self, '错误', "输入ip无效", QMessageBox.Close, QMessageBox.Close)
        elif QValidator.Acceptable.value != http_state.value:
            QMessageBox.critical(self, '错误', "输入http端口无效", QMessageBox.Close, QMessageBox.Close)
        elif QValidator.Acceptable.value != grpc_state.value:
            QMessageBox.critical(self, '错误', "输入grpc端口无效", QMessageBox.Close, QMessageBox.Close)
        elif server_name == "":
            QMessageBox.critical(self, '错误', "输入服务名称无效", QMessageBox.Close, QMessageBox.Close)
        elif server_name in self.parent_widget.server_info.keys():
            QMessageBox.critical(self, '错误', "输入服务名称重复, 删除旧的再新增", QMessageBox.Close, QMessageBox.Close)
        else:
            self.server_info["ip_addr"] = ip_addr
            self.server_info["http_port"] = http_port
            self.server_info["grpc_port"] = grpc_port
            self.server_info["server_name"] = server_name
            self.close()

    def cancle_button_callback(self):
        self.server_info = {}
        self.close()

class ServerManageWidget(QWidget):
    def __init__ (self, parent = None):
        super(ServerManageWidget, self).__init__(parent)
        config_layout = QVBoxLayout()
        self.server_info_widget = self.create_server_info_widget()
        self.output_info_widget = self.create_output_info_widget()
        config_layout.addWidget(self.server_info_widget)
        config_layout.addWidget(self.output_info_widget)
        self.setLayout(config_layout)
        self.server_info_dialog = ServerInfoDialog(self)
        self.server_info_dialog.setWindowTitle("服务信息")
        self.server_info = parent.server_info

    def create_output_info_widget(self):
        output_info_widget = QTextEdit()
        output_info_widget.setReadOnly(True)
        output_info_widget.setAutoFillBackground(True)
        output_info_widget.setStyleSheet("QTextEdit { font: 13px \"Courier New\";}")
        return output_info_widget

    def create_server_info_widget(self):
        server_info_widget = QWidget()
        server_info_layout = QHBoxLayout()
        self.server_list_widget = self.create_server_list_widget()
        button_group_widget = self.create_button_group_widget()
        server_info_layout.addWidget(self.server_list_widget)
        server_info_layout.addWidget(button_group_widget)
        server_info_widget.setLayout(server_info_layout)
        return server_info_widget

    def create_server_list_widget(self):
        server_list_widget = QListWidget()
        return server_list_widget

    def create_button_group_widget(self):
        button_group_widget = QWidget()
        button_group_layout = QGridLayout()
        btn_add = QPushButton("增加服务")
        btn_add.clicked.connect(self.button_add_clicked)
        btn_delete = QPushButton("删除服务")
        btn_delete.clicked.connect(self.button_delete_clicked)
        btn_status = QPushButton("服务状态")
        btn_status.clicked.connect(self.button_status_clicked)
        btn_models = QPushButton("模型列表")
        btn_models.clicked.connect(self.button_models_clicked)
        btn_statistic = QPushButton("服务统计")
        btn_statistic.clicked.connect(self.button_statistic_clicked)
        button_group_layout.addWidget(btn_add, 0, 0)
        button_group_layout.addWidget(btn_delete, 0, 1)
        button_group_layout.addWidget(btn_status, 1, 0)
        button_group_layout.addWidget(btn_models, 1, 1)
        button_group_layout.addWidget(btn_statistic, 2, 0)
        button_group_widget.setLayout(button_group_layout)
        return button_group_widget
    
    def get_all_server_info(self):
        item_num = self.server_list_widget.count()
        all_items = []
        for i in range(item_num):
            curr_item = self.server_list_widget.item(i)
            all_items.append(curr_item.text())
        return all_items
    
    def get_selected_server_info(self):
        select_items = self.server_list_widget.selectedItems()
        server_names = []
        for i in range(len(select_items)):
            server_names.append(select_items[i].text())
        if 0 == len(server_names):
            server_names = self.get_all_server_info()
        return server_names

    def button_add_clicked(self):
        self.server_info_dialog.clear_dialog()
        return_code = self.server_info_dialog.exec()
        server_info = self.server_info_dialog.server_info
        if server_info == {}:
            return

        server_name = server_info["server_name"]
        server_infos = self.get_all_server_info()
        if server_info["server_name"] not in self.server_info.keys():
            self.server_info[server_name] = server_info
        if server_name not in server_infos:
            item = QListWidgetItem(server_name)
            self.server_list_widget.addItem(item)

    def button_delete_clicked(self):
        select_items = self.server_list_widget.selectedItems()
        for i in range(len(select_items)):
            server_name = select_items[i].text()
            self.server_list_widget.takeItem(self.server_list_widget.row(select_items[i]))
            del self.server_info[server_name]

    def button_status_clicked(self):
        select_items = self.get_selected_server_info()
        for i in range(len(select_items)):
            server_name = select_items[i]
            if server_name in self.server_info.keys():
                ip_addr = self.server_info[server_name]["ip_addr"]
                http_port = self.server_info[server_name]["http_port"]
                grpc_port = self.server_info[server_name]["grpc_port"]
                client = TritonServerClient(ip_addr=ip_addr, http_port=http_port, grpc_port=grpc_port)
                try:
                    server_alive = client.is_server_alive()
                    server_ready = client.is_server_ready()
                    server_status_str = "service:{} ---- {}:{} alive status is {}, ready status is {}".format(server_name, 
                        ip_addr, grpc_port, server_alive, server_ready)
                except Exception as e:
                    server_status_str = "get server:{} ---- {}:{} status fail:\n{}".format(server_name,
                        ip_addr, grpc_port, e)
                self.output_info_widget.append(server_status_str)

    def button_models_clicked(self):
        select_items = self.get_selected_server_info()
        for i in range(len(select_items)):
            server_name = select_items[i]
            if server_name in self.server_info.keys():
                ip_addr = self.server_info[server_name]["ip_addr"]
                http_port = self.server_info[server_name]["http_port"]
                grpc_port = self.server_info[server_name]["grpc_port"]
                client = TritonServerClient(ip_addr=ip_addr, http_port=http_port, grpc_port=grpc_port)
                try:
                    server_statistic_info = client.get_inference_statistics()
                    server_models_infos = []
                    for model_info in server_statistic_info["model_stats"]:
                        model_info_dict = {}
                        model_name = model_info["name"]
                        model_version = model_info["version"]
                        model_metadata = client.get_model_metadata(model_name=model_name, model_version=model_version)
                        model_info_dict.update(model_metadata)
                        # model_config = client.get_model_config(model_name=model_name, model_version=model_version)
                        # model_info_dict.update(model_config)
                        server_models_infos.append(model_info_dict)
                    server_models_str = "service:{} ---- {}:{} models is\n{}".format(server_name, 
                        ip_addr, grpc_port, json.dumps(server_models_infos, indent=4))
                except Exception as e:
                    server_models_str = "get server:{} ---- {}:{} models fail:\n{}".format(server_name,
                        ip_addr, grpc_port, e)
                self.output_info_widget.append(server_models_str)

    def button_statistic_clicked(self):
        select_items = self.get_selected_server_info()
        for i in range(len(select_items)):
            server_name = select_items[i]
            if server_name in self.server_info.keys():
                ip_addr = self.server_info[server_name]["ip_addr"]
                http_port = self.server_info[server_name]["http_port"]
                grpc_port = self.server_info[server_name]["grpc_port"]
                client = TritonServerClient(ip_addr=ip_addr, http_port=http_port, grpc_port=grpc_port)
                try:
                    server_statistic_info = client.get_inference_statistics()
                    server_statistic_str = "get server:{} ---- {}:{} statistic:\n{}".format(server_name,
                        ip_addr, grpc_port, json.dumps(server_statistic_info, indent=4))
                except Exception as e:
                    server_statistic_str = "get server:{} ---- {}:{} statistic fail:\n{}".format(server_name,
                        ip_addr, grpc_port, e)
                self.output_info_widget.append(server_statistic_str)